package scu.maqiang.plasticity;

import scu.maqiang.fes.BVPType;
import scu.maqiang.fes.FES1L21;
import scu.maqiang.fes.FES2T31;
import scu.maqiang.mesh.*;
import scu.maqiang.numeric.*;

import java.util.Arrays;
import java.util.function.Function;

public class NonlinearSolve {
    public static void main(String[] args) {
        //DirectIteration();
        NewtonIteration_Case2();
        //NonlinearElasticity_IncreasedForce();
        //NonlinearElasticity_IncreasedDisplacement();
        //elastoplasticity_Case1();
    }

    //非线性泊松方程的直接迭代法
    public static void DirectIteration() {
        //1. 生成网格
        Mesh1L2 mesh1D = new Mesh1L2();
        mesh1D.line(0, 10, 10);

        //2.构造有限元空间
        FES1L21 fs = new FES1L21(mesh1D);
        SRMatrix A = new SRMatrix(fs.GetNdof());
        double[] RHS = new double[fs.GetNdof()];
        double[][] grad = new double[2][fs.GetNdof()];
        double[] coef = new double[fs.GetNdof()];
        Arrays.fill(coef, 1.0);
        //设置非线性系数
        double[][] NLCoef = {coef};
        NewIterSSolver solver = new NewIterSSolver();
        double[] x0 = new double[fs.GetNdof()];
        double[] x = new double[fs.GetNdof()];
        ScalarFunc exactFunc = (xy, label, param) -> 0.1 * (-1 + Math.sqrt(1 + 12  * xy[0]));
        double[] exactSol = mesh1D.functionInMesh(exactFunc, null);
        while(true) {
            fs.assembleStiff(NLCoef, BVPType.COMMON, A);
            fs.applyBC_MBN(A, RHS, Direct.All, 0.0, 1);
            fs.applyBC_MBN(A, RHS, Direct.All, 1.0, 2);
            solver.setMatrix(A);
            solver.PCGSSOR(RHS, x, 1.5, 1);
            if(MVO.distance(x0, x) < 1.0e-10) {
                break;
            }
            fs.computeGradient(x, grad);
            Arrays.setAll(x0, i -> x[i]);
            Arrays.setAll(coef, i -> 1.0 + 10 * x0[i]);
            A.clear();
            double[] reaction = fs.computeReactionForce(NLCoef, BVPType.COMMON, x);
            System.out.println("Reaction: ");
            System.out.println(Arrays.toString(reaction));
        }
        System.out.println("Numer Solution: ");
        System.out.println(Arrays.toString(x));
        System.out.println("Exact Solution: ");
        System.out.println(Arrays.toString(exactSol));
//        double[] reaction = fs.computeReactionForce(NLCoef, BVPType.COMMON, x);
//        System.out.println("Reaction: ");
//        System.out.println(Arrays.toString(reaction));
    }

    //求解非线性泊松方程的Newton迭代法
    public static void NewtonIteration_Case1() {
        Mesh1L2 mesh1D = new Mesh1L2();
        mesh1D.line(0, 10, 10);
        FES1L21 fs = new FES1L21(mesh1D);
        SRMatrix A = new SRMatrix(fs.GetNdof());
        double[] RHS = new double[fs.GetNdof()];
        double[] R = new double[fs.GetNdof()];
        double[][] grad = new double[2][fs.GetNdof()];
        //系数为a+bu = 1 + 10 * u
        double[] Acoef = MVO.vec(fs.GetNdof(), 10.0);
        double[] zeroCoef = new double[fs.GetNdof()];
        double[] Bcoef = MVO.vec(fs.GetNdof(), 100.0);
        double[] x = new double[fs.GetNdof()];
        double[][] NLCoef = {Acoef, zeroCoef, zeroCoef, zeroCoef, Bcoef, zeroCoef, zeroCoef, zeroCoef, zeroCoef, zeroCoef, x};
        //NONLINEAR_1用于组装系数为a + bx + cx^2 + dx^3 + e u + f u^2 + g u^3 + h du + i du^2 + j du^3 系数的矩阵
        fs.assembleStiff(NLCoef, BVPType.NONLINEAR_1, A);
        fs.applyBC_MBN(A, RHS, Direct.All, 0.0, 1);
        fs.applyBC_MBN(A, RHS, Direct.All, 1.0, 2);
        NewIterSSolver solver = new NewIterSSolver();
        solver.setMatrix(A);
        solver.PCGSSOR(RHS, x, 1.5, 1);
        double[] dx = new double[fs.GetNdof()];
        while(true) {
            A.clear();
            double[] reaction = fs.computeReactionForce(NLCoef, BVPType.NONLINEAR_1, x);
            Arrays.setAll(R, i -> RHS[i] - reaction[i]);
            //NONLINEAR_2用于组装系数为a + bx + cx^2 + dx^3 + e u + f u^2 + g u^3 + h du + i du^2 + j du^3 系数牛顿迭代的矩阵
            fs.assembleStiff(NLCoef, BVPType.NONLINEAR_2, A);
            fs.applyBC_MBN(A, R, Direct.All, 0.0, 1, 2);
            solver.setMatrix(A);
            solver.PCGSSSOR(R, dx, 1.5, 1);
            if(MVO.L2Norm(dx) < 1.0e-10) {
                break;
            }
            fs.computeGradient(x, grad);
            Arrays.setAll(x, i -> x[i] + dx[i]);

        }
        System.out.println("Solution: ");
        System.out.println(Arrays.toString(x));

    }

    //求解非线性泊松方程的Newton迭代法
    public static void NewtonIteration_Case2() {
        Mesh1L2 mesh1D = new Mesh1L2();
        mesh1D.line(0, 10, 10);
        FES1L21 fs = new FES1L21(mesh1D);
        SRMatrix A = new SRMatrix(fs.GetNdof());
        double[] RHS = new double[fs.GetNdof()];
        double[] R = new double[fs.GetNdof()];
        //系数为a+bu = 10 + 100 * u
        double[] Acoef = MVO.vec(fs.GetNdof(), 10.0);
        double[] zeroCoef = new double[fs.GetNdof()];
        double[] Bcoef = MVO.vec(fs.GetNdof(), 100.0);
        NewIterSSolver solver = new NewIterSSolver();
        double[] x = new double[fs.GetNdof()];
        double[][] NLCoef = {Acoef, zeroCoef, zeroCoef, zeroCoef, Bcoef, zeroCoef, zeroCoef, zeroCoef, zeroCoef, zeroCoef, x};
        fs.assembleStiff(NLCoef, BVPType.NONLINEAR_1, A);
        fs.applyBC_MBN(A, RHS, Direct.All, 0.0, 1);
        fs.assembleSource(new double[]{0.2}, BVPType.COMMON, RHS);
        System.out.println(Arrays.toString(RHS));
        solver.setMatrix(A);
        solver.PCGSSOR(RHS, x, 1.5, 1);
        double[] dx = new double[fs.GetNdof()];
        ScalarFunc exactFunc = (xy, label, param) -> 0.1 * (-1 + Math.sqrt(1 + 0.2  * xy[0] * (20 - xy[0])));
        double[] exactSol = mesh1D.functionInMesh(exactFunc, null);
        while(true) {
            A.clear();
            double[] reaction = fs.computeReactionForce(NLCoef, BVPType.NONLINEAR_1, x);
            Arrays.setAll(R, i -> RHS[i] - reaction[i]);
            fs.assembleStiff(NLCoef, BVPType.NONLINEAR_2, A);
            fs.applyBC_MBN(A, R, Direct.All, 0.0, 1);
            solver.setMatrix(A);
            solver.PCGSSSOR(R, dx, 1.5, 1);
            if(MVO.L2Norm(dx) < 1.0e-10) {
                break;
            }
            Arrays.setAll(x, i -> x[i] + dx[i]);
        }
        System.out.println("Numer Solution: ");
        System.out.println(Arrays.toString(x));
        System.out.println("Exact Solution: ");
        System.out.println(Arrays.toString(exactSol));
    }

    //求解非线性弹性方程的载荷增量法
    public static void NonlinearElasticity_IncreasedForce() {
        Mesh1L2 mesh1D = new Mesh1L2();
        mesh1D.line(0, 10, 2);
        FES1L21 fs = new FES1L21(mesh1D);
        SRMatrix A = new SRMatrix(fs.GetNdof());
        double[] RHS = new double[fs.GetNdof()];
        double[] R = new double[fs.GetNdof()];
        //应变函数导数E0(1 - 10 * du), E0 = 200
        double[] Acoef = new double[fs.GetNdof()];
        Arrays.fill(Acoef, 200.0);
        double[] zeroCoef = new double[fs.GetNdof()];
        double[] Bcoef = new double[fs.GetNdof()];
        Arrays.fill(Bcoef, -2000.0);
        double[] x = new double[fs.GetNdof()];
        double[] dx = new double[fs.GetNdof()];
        double[][] NLCoef = {Acoef, zeroCoef, zeroCoef, zeroCoef, zeroCoef, zeroCoef, zeroCoef, Bcoef, zeroCoef, zeroCoef, x};
        double[] force = {8.0, 2.0};
        Arrays.fill(Acoef, 200.0);
        Arrays.fill(Bcoef, -1000.0);
        NewIterSSolver solver = new NewIterSSolver();
        //应变函数E0(du - 5 * du * du)
        Function<Double, Double> forceFunc = strain -> 200 * (strain - 5 * strain * strain);
        for(int i = 0; i < force.length; i++) {
            fs.assembleStiff(NLCoef, BVPType.NONLINEAR_1, A);
            fs.assembleFlux(new double[]{force[i]}, BVPType.COMMON, RHS, 2);
            System.out.println(Arrays.toString(RHS));
            fs.applyBC_MBN(A, RHS, Direct.All, 0.0, 1);
            solver.setMatrix(A);
            solver.PCGSSOR(RHS, x, 1.5, 1);
            while(true) {
                double[] reaction = fs.computeReactionForce(forceFunc, BVPType.COMMON, x);
                //System.out.println(Arrays.toString(reaction));
                System.out.println("Action Force: " + Arrays.toString(reaction));
                Arrays.setAll(R, k -> RHS[k] - reaction[k]);
                fs.applyBC_MBN(R, Direct.All, 0.0, 1);
                System.out.println("Residual Force: " + Arrays.toString(R));
                solver.PCGSSOR(R, dx, 1.5, 0);
                if(MVO.L2Norm(dx) < 1.0e-10) {
                    break;
                }
                Arrays.setAll(x, k -> x[k] + dx[k]);
                //System.out.println(Arrays.toString(x));
                //System.exit(1);
            }
            //Arrays.setAll(RHS, k -> R[k]);
            System.out.println(Arrays.toString(RHS));
            A.clear();
            System.out.println(Arrays.toString(x));
        }
    }

    //求解非线性弹性方程的位移增量法
    public static void NonlinearElasticity_IncreasedDisplacement() {
        Mesh1L2 mesh1D = new Mesh1L2();
        mesh1D.line(0, 10, 2);
        FES1L21 fs = new FES1L21(mesh1D);
        SRMatrix A = new SRMatrix(fs.GetNdof());
        double[] RHS = new double[fs.GetNdof()];
        double[] R = new double[fs.GetNdof()];
        //应变函数导数E0(1 - 10 * du), E0 = 200
        double[] Acoef = new double[fs.GetNdof()];
        Arrays.fill(Acoef, 200.0);
        double[] zeroCoef = new double[fs.GetNdof()];
        double[] Bcoef = new double[fs.GetNdof()];
        Arrays.fill(Bcoef, -2000.0);
        double[] x = new double[fs.GetNdof()];
        double[] dx = new double[fs.GetNdof()];
        double[][] NLCoef = {Acoef, zeroCoef, zeroCoef, zeroCoef, zeroCoef, zeroCoef, zeroCoef, Bcoef, zeroCoef, zeroCoef, x};
        //设置位移增量
        double[] displacement = {0.6, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2};
        Arrays.fill(Acoef, 200.0);
        Arrays.fill(Bcoef, -1000.0);
        NewIterSSolver solver = new NewIterSSolver();
        //应变函数E0(du - 5 * du * du)
        Function<Double, Double> forceFunc = strain -> 200 * (strain - 5 * strain * strain);
        double totalDisp = 0.0;
        fs.assembleStiff(NLCoef, BVPType.NONLINEAR_1, A);
        fs.applyBC_MBN(A, Direct.All, 1);
        fs.applyBC_MBN(A, Direct.All, 2);
        solver.setMatrix(A);
        for(int i = 0; i < displacement.length; i++) {
            totalDisp += displacement[i];
            fs.applyBC_MBN(RHS, Direct.All, 0.0, 1);
            fs.applyBC_MBN(RHS, Direct.All, totalDisp, 2);
            solver.PCGSSOR(RHS, x, 1.5, 1);
            while (true) {
                double[] reaction = fs.computeReactionForce(forceFunc, BVPType.COMMON, x);
                //System.out.println(Arrays.toString(reaction));
                System.out.println("Reaction Force: " + Arrays.toString(reaction));
                Arrays.setAll(R, k -> RHS[k] - reaction[k]);
                fs.applyBC_MBN(R, Direct.All, 0.0, 1, 2);
                System.out.println("Residual Force: " + Arrays.toString(R));
//                System.out.println(A.display());
                solver.PCGSSOR(R, dx, 1.5, 1);
                System.out.println("Increased displacement: " + Arrays.toString(dx));
                if (MVO.L2Norm(dx) < 1.0e-7) {
                    break;
                }
                System.out.println("Total displacement: " + Arrays.toString(dx));
                Arrays.setAll(x, k -> x[k] + dx[k]);
                //System.out.println(Arrays.toString(x));
                //System.exit(1);
            }
            System.out.println("Final displacement: " + Arrays.toString(x));
        }
    }

    //使用初始刚度法求解线塑性问题
    public static void elastoplasticity_Case1() {
        String meshFile = "meshPlasticity1D.msh";
        Mesh1L2 mesh = new Mesh1L2();
        mesh.readMesh(MeshFormat.msh, meshFile);
        FES1L21 fs = new FES1L21(mesh);
        SRMatrix A = new SRMatrix(fs.GetNdof());
        double[] RHS = new double[fs.GetNdof()];
        double[] R = new double[fs.GetNdof()];
        double[] x = new double[fs.GetNdof()];
        double[] dx = new double[fs.GetNdof()];
        ScalarFunc elasticModuleFunc = (xy, label, param) -> 10000;
        ScalarFunc areaFunc = (xy, label, param) -> label ==1?1.0:2.0;
        ScalarFunc yieldFunc = (xy, label, param) -> label == 1?5.0: 7.5;
        ScalarFunc hardFunc = (xy, label, param) -> label == 1?1000:2000;

        //刚度系数, 等于弹性模量乘以横截面积
        ScalarFunc EFunc = (xy, label, param) -> label == 1?10000:20000;
        //塑性系数, 包含4个参数函数: 弹性模量, 横截面积, 屈服应力, 强化系数
        ScalarFunc[] plasticCoef = {elasticModuleFunc, areaFunc, yieldFunc, hardFunc};

        //组装初始弹性刚度矩阵
        fs.assembleStiff(new ScalarFunc[]{EFunc}, null, BVPType.COMMON, A);
        fs.applyBC_MBN(A, Direct.All, 1);

        NewIterSSolver solver = new NewIterSSolver();
        solver.setMatrix(A);

        //载荷增量向量
        double[] force = {12.5, 2.5, 2.5, 2.5, 2.5, 2.5,
                          -35, -2.5, -2.5, -2.5, -2.5, -2.5, -2.5, -2.5, -2.5, -2.5,
                          5.0, 5.0, 5.0, 5.0, 5.0, 5.0, 5.0, 5.0, 5.0, 5.0, 5.0, 5.0,
                          -10, -10, -10, -2,5, -2.5, -2.5, -2,5, -2.5, -2.5, -5.0, -5.0, -5.0, -5.0, -5.0};
        //支反力向量
        double[] reaction = new double[fs.GetNdof()];
        double[] plotX = new double[force.length + 1];
        double[] plotY = new double[force.length + 1];
        for(int i = 0; i < force.length; i++) {
            Arrays.fill(dx, 0.0);
            fs.assembleFlux(new double[]{force[i]}, BVPType.COMMON, RHS, 2);

            while (true) {
                double[] rr = reaction;
                Arrays.setAll(R, k -> RHS[k] - rr[k]);
//              System.out.println("Residual Force: " + R[5]);
                fs.applyBC_MBN(R, Direct.All, 0.0, 1);
                solver.PCGSSOR(R, dx, 1.5, 0);
                //Pair<double[], double[]> strainInfo = fs.computeElementGradient(dx);
//                System.out.println("dx: " + dx[5]);
                if (MVO.L2Norm(dx) < 1.0e-10) {
                    break;
                }
                Arrays.setAll(x, k -> x[k] + dx[k]);
                reaction = fs.computeElastoplasticActionForce(plasticCoef, null, BVPType.COMMON, dx);
                //System.out.println("Action Force: " + Arrays.toString(reaction));
                //System.out.println("Solution x:" + Arrays.toString(x));
                //System.out.println("===============================");
                //System.exit(1);
            }
//            System.out.println(RHS[5] + "\t " + x[5]);
            plotX[i + 1] = x[5];
            plotY[i + 1] = plotY[i] + force[i];
        }
        Tecplot.LineXY("Load-Disp.dat", plotX, new double[][]{plotY});
    }

    public static void elastoplasticity_Case2() {

    }
}
